package utils;

import com.mongodb.MongoClient;
import com.mongodb.client.MongoCursor;
import org.bson.Document;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;

/**
 * Created by myc on 2021/4/22.
 */
public class MongoMultiCursor implements Iterator<List<Document>> {

    private List<MongoCursorBean> mongoCursorBeans;

    private Document projection;

    public MongoMultiCursor(MongoClient mongoClient, List<String> tables, String db) {

        mongoCursorBeans = new ArrayList<>(tables.size());
        for (String table : tables) {
            MongoCursorBean mongoCursorBean = new MongoCursorBean();
            mongoCursorBean.table = table;
            mongoCursorBean.mongoCursor = mongoClient.getDatabase(db).getCollection(table).find().sort(new Document("_id", 1)).iterator();
            mongoCursorBeans.add(mongoCursorBean);
        }
    }

    public MongoMultiCursor setProjection(Document projection) {
        this.projection = projection;
        return this;
    }

    @Override
    public boolean hasNext() {
        for (MongoCursorBean mongoCursorBean : mongoCursorBeans) {
            if (mongoCursorBean.hasNext()){
                return true;
            }
        }
        return false;
    }

    @Override
    public List<Document> next() {
        return getMinDocuments();
    }

    private List<Document> getMinDocuments(){
        List<Document> ret = new ArrayList<>(mongoCursorBeans.size());
        List<SortBean> sortBeans = new ArrayList<>();
        for (int i = 0; i < mongoCursorBeans.size(); i++) {
            Document head = mongoCursorBeans.get(i).head();
            if (head != null){
                sortBeans.add(new SortBean(i, head.getString("_id")));
            }
        }
        sortBeans.sort(Comparator.comparing(sortBean -> sortBean.id));

        String currentId = null;
        for (SortBean sortBean : sortBeans) {
            if (currentId == null){
                currentId = sortBean.id;
                MongoCursorBean mongoCursorBean = mongoCursorBeans.get(sortBean.idx);
                if (mongoCursorBean.hasNext()){
                    ret.add(mongoCursorBean.next());
                }
            }else {
                if (currentId.equals(sortBean.id)){
                    MongoCursorBean mongoCursorBean = mongoCursorBeans.get(sortBean.idx);
                    if (mongoCursorBean.hasNext()){
                        ret.add(mongoCursorBean.next());
                    }
                }else {
                    break;
                }
            }
        }
        return ret;
    }

    private static class MongoCursorBean{
        private String table;
        private Document head;
        private MongoCursor<Document> mongoCursor;

        public boolean hasNext() {
            return mongoCursor.hasNext() || head != null;
        }

        public Document next() {
            if (head == null){
                head = mongoCursor.hasNext()? mongoCursor.next(): null;
                return head;
            }else {
                Document ret = head;
                head = mongoCursor.hasNext()? mongoCursor.next(): null;
                return ret;
            }

        }
        public Document head(){
            if (head == null){
                head = mongoCursor.hasNext()? mongoCursor.next(): null;
            }
            return head;
        }
    }

    private static class SortBean{
        private int idx;
        private String id;

        private SortBean(int idx, String id) {
            this.idx = idx;
            this.id = id;
        }
    }
}
